package belf.migrate.training;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * 大模型和数据库对应关系
 */
public class LlmDB {
    // 保存可调用的大模型的名称
    public static final List<String> LLMs = new ArrayList<>();
    // 保存大模型名称到数据库名称（或数据库schema名称）的映射
    public static final Map<String, String> LLM2DBs = new HashMap<>();

    static {
        LLMs.add("Doubao-pro-32k");
        LLMs.add("ERNIE-4.0-8K");
        LLMs.add("GLM-4-Plus");
//        LLMs.add("gemma2:27b");
        LLMs.add("gpt-4o");
        LLMs.add("qwen-max");
        LLMs.add("hunyuan-large");

        LLM2DBs.put("Doubao-pro-32k", "Doubao");
        LLM2DBs.put("ERNIE-4.0-8K", "ERNIE");
        LLM2DBs.put("GLM-4-Plus", "GLM");
//        LLM2DBs.put("gemma2:27b", "gemma2");
        LLM2DBs.put("gpt-4o", "gpt4o");
        LLM2DBs.put("qwen-max", "qwen-max");
        LLM2DBs.put("hunyuan-large", "hunyuan-large");

        // TODO:需要再加上星火大模型、商汤大模型、零一万物大模型
    }
}
